/**
 * @file IsReal_custom_tiling.cpp
 *
 * Copyright (C) 2025. Huawei Technologies Co., Ltd. All rights reserved.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
 */
#include <cassert>
#include <cstddef>
#include "is_real_custom_tiling.h"

constexpr uint32_t DATA_TYPE_SIZE[] = {
    2,   // 0: IsReal_BFLOAT16 (bfloat16)
    2,   // 1: IsReal_FLOAT16 (float16)
    4,   // 2: IsReal_FLOAT32 (float32)
    1,   // 3: IsReal_INT8 (int8)
    2,   // 4: IsReal_INT16 (int16)
    4,   // 5: IsReal_INT32 (int32)
    1,   // 6: IsReal_UINT8 (uint8)
    2,   // 7: IsReal_UINT16 (uint16)
    4,   // 8: IsReal_UINT32 (uint32)
    8,   // 9: IsReal_INT64 (int64)
    8,   // 10: IsReal_UINT64 (uint64)
    1,   // 11: IsReal_BOOL (bool)
    8,   // 12: IsReal_DOUBLE (double/float64)
    0,   // 13: IsReal_STRING (string，可变长度)
    1,   // 14: IsReal_DUAL_SUB_INT8 (dual int8)
    1,   // 15: IsReal_DUAL_SUB_UINT8 (dual uint8)
    8,   // 16: IsReal_COMPLEX64 (2*float32)
    16,  // 17: IsReal_COMPLEX128 (2*double)
    1,   // 18: IsReal_QINT8 (量化int8)
    2,   // 19: IsReal_QINT16 (量化int16)
    4,   // 20: IsReal_QINT32 (量化int32)
    1,   // 21: IsReal_QUINT8 (量化uint8)
    2,   // 22: IsReal_QUINT16 (量化uint16)
    0,   // 23: IsReal_RESOURCE (资源类型，大小不固定)
    0,   // 24: IsReal_STRING_REF (字符串引用)
    0,   // 25: IsReal_DUAL (双输出类型)
    0,   // 26: IsReal_VARIANT (变体类型)
    0,   // 27: IsReal_UNDEFINED (未定义类型)
    1,   // 28: IsReal_INT4 (int4，按1字节存储)
    1,   // 29: IsReal_UINT1 (uint1，按1字节存储)
    1,   // 30: IsReal_INT2 (int2，按1字节存储)
    1,   // 31: IsReal_UINT2 (uint2，按1字节存储)
    4,   // 32: IsReal_COMPLEX32 (2*float16)
    1,   // 33: IsReal_HIFLOAT8 (float8)
    1,   // 34: IsReal_FLOAT8_E5M2 (float8)
    1,   // 35: IsReal_FLOAT8_E4M3FN (float8)
    1,   // 36: IsReal_FLOAT8_E8M0 (float8)
    1,   // 37: IsReal_FLOAT6_E3M2 (float6，按1字节存储)
    1,   // 38: IsReal_FLOAT6_E2M3 (float6，按1字节存储)
    1,   // 39: IsReal_FLOAT4_E2M1 (float4，按1字节存储)
    1,   // 40: IsReal_FLOAT4_E1M2 (float4，按1字节存储)
    0    // 41: IsReal_MAX (边界标记)
};
constexpr uint32_t BLOCK_SIZE = 32;
constexpr uint32_t BUFFER_NUM = 2;
constexpr uint32_t UB_BLOCK_NUM = 100;  // UB最大可以使用的block数量
constexpr uint32_t MAX_AVAILABLE_UB_BLOCK_NUM = UB_BLOCK_NUM / BUFFER_NUM * BUFFER_NUM;

// tiling参数计算函数
void TilingParamsCalc(uint32_t length, uint32_t alignNum,
    uint32_t& tileNum, uint32_t& tileLength, uint32_t& lastTileLength)
{
    assert(alignNum != 0U);
    tileNum = length / (alignNum * MAX_AVAILABLE_UB_BLOCK_NUM);
    if ((static_cast<uint32_t>(length / alignNum) % MAX_AVAILABLE_UB_BLOCK_NUM == 0U) || tileNum == 0U) {
        if (tileNum == 0U) {
            tileNum = 1U;
        }
        if (length < MAX_AVAILABLE_UB_BLOCK_NUM * alignNum) {
            tileLength = ((static_cast<size_t>(length) + alignNum - 1) / alignNum) * alignNum;
            lastTileLength = tileLength;
        } else {
            tileLength = MAX_AVAILABLE_UB_BLOCK_NUM * alignNum;
            lastTileLength = (uint32_t)(length - (tileNum - 1) * tileLength);
        }
    } else {
        tileNum++;
        tileLength = MAX_AVAILABLE_UB_BLOCK_NUM * alignNum;
        lastTileLength = (uint32_t)(length - (tileNum - 1) * tileLength);
    }
}

void GenerateTilingData(uint8_t* tilingBuf, uint32_t blockDim)
{
    uint32_t totalLength;
    uint32_t dataTypeSize;
    uint32_t blockLength;
    uint32_t totalLengthAligned;

    IsRealCustomTilingData *tiling = reinterpret_cast<IsRealCustomTilingData *>(tilingBuf);
    totalLength = tiling->totalLength;
    dataTypeSize = DATA_TYPE_SIZE[tiling->dataType];

    uint32_t alignNum = BLOCK_SIZE / dataTypeSize;
    assert((alignNum != 0U) && (blockDim != 0U));
    /** 计算使用的核数 **/
    /* 如果传入数据的长度非32B对齐, 计算对齐后的长度*/
    totalLengthAligned = (totalLength % alignNum == 0U)?
        static_cast<uint32_t>(totalLength) :
        ((static_cast<uint32_t>(totalLength) + alignNum - 1) / alignNum) * alignNum;

    /* 核间可均分场景 */
    if ((totalLengthAligned / alignNum) % blockDim == 0U) {
        uint32_t tileNum = 0;
        uint32_t tileLength = 0;
        uint32_t lastTileLength = 0;
        blockLength = totalLengthAligned / blockDim;
        TilingParamsCalc(blockLength, alignNum, tileNum, tileLength, lastTileLength);

        tiling->blockLength = blockLength;
        tiling->tileNum = tileNum;
        tiling->tileLength = tileLength;
        tiling->lastTileLength = lastTileLength;
        tiling->isEvenCore = 1U;
    } else {  // 核间不可均分
        uint32_t formerNum = (totalLengthAligned / alignNum) % blockDim;
        uint32_t tailNum = blockDim - formerNum;
        // 计算整块和尾块的数据量
        uint32_t formerLength =
            static_cast<uint32_t>(((totalLengthAligned + blockDim - 1) / blockDim + alignNum - 1) / alignNum) * alignNum;
        uint32_t tailLength = (totalLengthAligned / blockDim / alignNum) * alignNum;

        uint32_t formerTileNum;
        uint32_t formerTileLength;
        uint32_t formerLastTileLength;

        uint32_t tailTileNum;
        uint32_t tailTileLength;
        uint32_t tailLastTileLength;

        TilingParamsCalc(formerLength, alignNum,
            formerTileNum, formerTileLength, formerLastTileLength);
        TilingParamsCalc(tailLength, alignNum,
            tailTileNum, tailTileLength, tailLastTileLength);

        tiling->formerNum = formerNum;
        tiling->formerLength = formerLength;
        tiling->formerTileNum = formerTileNum;
        tiling->formerTileLength = formerTileLength;
        tiling->formerLastTileLength = formerLastTileLength;

        tiling->tailNum = tailNum;
        tiling->tailLength = tailLength;
        tiling->tailTileNum = tailTileNum;
        tiling->tailTileLength = tailTileLength;
        tiling->tailLastTileLength = tailLastTileLength;
        tiling->isEvenCore = 0U;
    }
}